# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0

import math
from isaaclab.utils import configclass

# 导入粗糙地形运动速度环境配置基类
from rl_training.tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
# # 可选：导入射线传感器的网格模式配置
##
# Pre-defined configs
##
# 导入机器人配置
from rl_training.assets.deeprobotics import DEEPROBOTICS_DUCK_CFG  # isort: skip
import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp

@configclass
class DuckRoughEnvCfg(LocomotionVelocityRoughEnvCfg):
    base_link_name = "base_link"
    foot_link_name = ".*_ankle"
    # fmt: off
    # joint_names = [
    #     "left_hip_pitch_joint",          # 0  L_LEG_HIP_PITCH
    #     "left_hip_roll_joint",           # 1  L_LEG_HIP_ROLL
    #     "left_hip_yaw_joint",            # 2  L_LEG_HIP_YAW
    #     "left_knee_joint",               # 3  L_LEG_KNEE
    #     "left_ankle_joint",              # 4  L_LEG_ANKLE
    #     "right_hip_pitch_joint",         # 5  R_LEG_HIP_PITCH
    #     "right_hip_roll_joint",          # 6  R_LEG_HIP_ROLL
    #     "right_hip_yaw_joint",           # 7  R_LEG_HIP_YAW
    #     "right_knee_joint",              # 8  R_LEG_KNEE
    #     "right_ankle_joint",             # 9  R_LEG_ANKLE
    # ]
    # link_names = [  # 机器人所有链接名称
    #     "left_hip_pitch",          # 0  L_LEG_HIP_PITCH
    #     "left_hip_roll",           # 1  L_LEG_HIP_ROLL
    #     "left_hip_yaw",            # 2  L_LEG_HIP_YAW
    #     "left_knee",               # 3  L_LEG_KNEE
    #     "left_ankle",              # 4  L_LEG_ANKLE
    #     "right_hip_pitch",         # 5  R_LEG_HIP_PITCH
    #     "right_hip_roll",          # 6  R_LEG_HIP_ROLL
    #     "right_hip_yaw",           # 7  R_LEG_HIP_YAW
    #     "right_knee",              # 8  R_LEG_KNEE
    #     "right_ankle",             # 9  R_LEG_ANKLE
    #     "base_link",               # 10 WAIST_YAW
    # ]
    # fmt: on

    def __post_init__(self):
        # post init of parent
        super().__post_init__()

        # ------------------------------Sence------------------------------
        self.scene.robot = DEEPROBOTICS_DUCK_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")   # 设置机器人模型路径
        self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name  # 高度扫描仪主路径
        self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name  # 高度扫描仪基准路径
        # self.scene.height_scanner.pattern_cfg.resolution = 0.07 # 设置高度扫描仪分辨率

        # ------------------------------Observations------------------------------
        self.observations.policy.base_lin_vel.scale = 0.25
        self.observations.policy.base_ang_vel.scale = 0.25  # 基础角速度缩放
        self.observations.policy.joint_pos.scale = 0.25     # 关节位置缩放
        self.observations.policy.joint_vel.scale = 0.05     # 关节速度缩放
        self.observations.policy.base_lin_vel = None        # 不观测基础线速度
        self.observations.policy.height_scan = None         # 不观测高度扫描
        # self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names  # 关节位置观测关节名
        # self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names  # 关节速度观测关节名

        # ------------------------------Actions------------------------------
        # reduce action scale 降低动作缩放比例
        self.actions.joint_pos.scale = 0.25
        self.actions.joint_pos.clip = {".*": (-100.0, 100.0)}  # 所有关节动作裁剪范围
        # self.actions.joint_pos.joint_names = self.joint_names  # 动作关节名

        # ------------------------------Events------------------------------
        # 随机重置基座参数
        # self.events.randomize_reset_base.params = {
        #     "pose_range": {
        #         "x": (-1.0, 1.0),
        #         "y": (-1.0, 1.0),
        #         "z": (0.0, 0.0),
        #         "roll": (-0.3, 0.3),
        #         "pitch": (-0.3, 0.3),
        #         "yaw": (-3.14, 3.14),
        #     },
        #     "velocity_range": {
        #         "x": (-0.2, 0.2),
        #         "y": (-0.2, 0.2),
        #         "z": (-0.2, 0.2),
        #         "roll": (-0.05, 0.05),
        #         "pitch": (-0.05, 0.05),
        #         "yaw": (-0.0, 0.0),
        #     },
        # }
        self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = [self.base_link_name] # 随机刚体质量的链接名
        self.events.randomize_rigid_body_mass_base = None  # 不随机基座质量
        self.events.randomize_com_positions.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_apply_external_force_torque.params["asset_cfg"].body_names = [self.base_link_name]

        # ------------------------------Rewards------------------------------
        # General
        self.rewards.is_terminated.weight = -200.0

        # Root penalties
        # Z轴方向线速度和XY方向角速度惩罚，不希望机器人在Z轴上下运动，或者在XY平面有过多旋转，期望平稳移动
        self.rewards.lin_vel_z_l2.weight = 0      # z方向线速度惩罚
        self.rewards.ang_vel_xy_l2.weight = -0.1  # xy方向角速度惩罚
        # 基座在Z轴方向保持平稳，不希望它翻滚或侧翻
        self.rewards.flat_orientation_l2.weight = -0.2 # 姿态平坦度惩罚
        self.rewards.base_height_l2.weight = 0.1       # 基座高度惩罚
        self.rewards.base_height_l2.params["target_height"] = 0.35 # 目标高度
        self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name] 
        self.rewards.body_lin_acc_l2.weight = 0
        self.rewards.body_lin_acc_l2.params["asset_cfg"].body_names = [self.base_link_name]

        # Joint penalties
        self.rewards.joint_torques_l2.weight = -1.5e-7 # 关节力矩惩罚
        self.rewards.joint_torques_l2.params["asset_cfg"].joint_names = [
            "left_hip_yaw_joint", "right_hip_yaw_joint",
            "left_hip_roll_joint", "right_hip_roll_joint",
            "left_hip_pitch_joint", "right_hip_pitch_joint",
            "left_knee_joint", "right_knee_joint",
            "left_ankle_joint", "right_ankle_joint"
        ]
        self.rewards.joint_vel_l2.weight = 0
        self.rewards.joint_acc_l2.weight = -1.25e-7
        self.rewards.joint_acc_l2.params["asset_cfg"].joint_names = [".*_hip_.*", ".*_knee_joint"]
        self.rewards.joint_pos_limits.weight = -0.5
        self.rewards.joint_vel_limits.weight = 0
        self.rewards.joint_power.weight = 0
        self.rewards.stand_still_without_cmd.weight = 0
        # 关节位置惩罚（静止时加重）
        self.rewards.joint_pos_penalty.weight = -1.0
        self.rewards.joint_mirror.weight = 0.05
        self.rewards.joint_mirror.params["mirror_joints"] = [
            ["left_hip_yaw_joint", "right_hip_yaw_joint"],
            ["left_hip_roll_joint", "right_hip_roll_joint"],
            ["left_hip_pitch_joint", "right_hip_pitch_joint"],
            ["left_knee_joint", "right_knee_joint"],
            ["left_ankle_joint", "right_ankle_joint"]
        ]

        # Action penalties
        self.rewards.action_rate_l2.weight = -0.02 # 动作变化率惩罚
        self.rewards.action_mirror.weight = 0.1
        self.rewards.action_mirror.params["mirror_joints"] = [
            ["left_hip_yaw_joint", "right_hip_yaw_joint"],
            ["left_hip_roll_joint", "right_hip_roll_joint"],
            ["left_hip_pitch_joint", "right_hip_pitch_joint"],
            ["left_knee_joint", "right_knee_joint"],
            ["left_ankle_joint", "right_ankle_joint"]
        ]

        # Contact sensor不期望它在地上翻滚
        self.rewards.undesired_contacts.weight = 0 # 非期望接触惩罚
        self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"]
        # 希望机器人足端轻触地面，接触力越小越好。
        self.rewards.contact_forces.weight = -2e-2 # 接触力惩罚
        self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name]

        # Velocity-tracking rewards
        self.rewards.track_lin_vel_xy_exp.weight = 1.2  # xy方向速度跟踪奖励
        self.rewards.track_lin_vel_xy_exp.func = mdp.track_lin_vel_xy_yaw_frame_exp
        self.rewards.track_ang_vel_z_exp.weight = 0.6   # z方向角速度跟踪奖励
        self.rewards.track_ang_vel_z_exp.func = mdp.track_ang_vel_z_world_exp
        # Others
        self.rewards.feet_air_time.weight = 2.0        # 足部离地时间奖励
        self.rewards.feet_air_time.func = mdp.feet_air_time_positive_biped
        self.rewards.feet_air_time.params["threshold"] = 0.6  # 离地阈值
        self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_contact.weight = 0
        self.rewards.feet_contact.params["sensor_cfg"].body_names = [self.foot_link_name]
        # 未收到指令时足部接触奖励
        self.rewards.feet_contact_without_cmd.weight = 0.05
        self.rewards.feet_contact_without_cmd.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_stumble.weight = 0
        self.rewards.feet_stumble.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.weight = -0.5   # 足部滑动惩罚
        self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name]
        # 踏步时，实际想要的足端高度
        self.rewards.feet_height.weight = -0.2  # 足部高度惩罚
        self.rewards.feet_height.params["target_height"] = 0.05
        self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_height_body.weight = 0  # 足部相对身体高度惩罚
        self.rewards.feet_height_body.params["target_height"] = -0.1
        self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_gait.weight = 0
        self.rewards.feet_gait.params["synced_feet_pair_names"] = (("left_ankle", "right_ankle"),)
        self.rewards.upward.weight = 1.0

        # If the weight of rewards is 0, set rewards to None
        # 如果奖励权重为0，则禁用该奖励
        if self.__class__.__name__ == "DuckRoughEnvCfg":
            self.disable_zero_weight_rewards()

        # ------------------------------Terminations------------------------------
        self.terminations.illegal_contact = None      # 非法接触终止条件
        # self.terminations.bad_orientation_2 = None  # 可选：姿态异常终止条件

        # ------------------------------Curriculums------------------------------
        # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0)  # 可选：命令难度课程
        self.curriculum.command_levels = None  # 不使用命令课程

        # ------------------------------Commands------------------------------
        self.commands.base_velocity.ranges.lin_vel_x = (-0.5, 0.75)  # x方向速度命令范围
        self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5)   # y方向速度命令范围
        self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0)   # z方向角速度命令范围